# -*- coding: utf-8 -*-
"""
Created on Wed Oct 19 16:11:35 2022

@author: ZJH
"""

from cion.core.hardware.rtmq.RTMQ_Assembler import BaseCore
import contextvars


# --------- Standard Core Structure ---------


class StdCore(BaseCore):
    """

    Standard RTMQ core structure
    with basic computation and flow control capability.

    Constants
    ----------
    RSV_GPR : Assembler reserved general purpose registers. \n
    N_PLM : Cache pipeline latency. \n
    N_PLI : Type-I instruction pipeline latency. \n
    N_PLA : Type-A instruction pipeline latency. \n
    N_PLS : Stack pop pipeline latency.


    """

    RSV_GPR = ["TR0", "TR1", "NOW",
               "AR0", "AR1", "AR2", "AR3"]

    N_PLM = 3
    N_PLI = 1
    N_PLA = 4
    N_PLS = 4

    def __init__(self, n_gpr, n_slv, regs):
        """

        Parameters
        ----------
        n_gpr : integer
            Number of non-reserved general-purpose registers.
        n_slv : integer
            Number of RTLink slaves (within 0-15).
        regs : list of strings
            List of all other registers.

        Returns
        -------
        None.

        """
        tmp = self.RSV_GPR
        for i in range(n_gpr):
            tmp += [f"R{i:X}"]
        for i in range(1, n_slv + 1):
            tmp += [f"RTD{i:X}"]
        tmp += regs
        super().__init__(tmp)


# --------- Compiler Class ---------


class Compiler:
    """

    A helper class used to handle control structures
    (if-else, for, while, etc.)
    and generate assembly program.

    Constants
    ----------
    MAX_DPT : Maximum depth of nested control structures.

    Methods
    ----------
    new_context() : initialize a new context for a control structure. \n
    pack_context() : pack current context and return the assembly code
    within the context. \n
    append_instr(ins) : append instructions to current context. \n
    new_subroutine(name, ins) : define a new subroutine. \n
    assemble() : assemble the program and subroutine definitions.

    """

    MAX_DPT = 50

    def __init__(self, core):
        self.stk = [[] for i in range(self.MAX_DPT)]
        self.core = core
        self.ptr = 0
        self.sub = {}
        self.asm = ""

    def new_context(self):
        self.ptr += 1
        self.stk[self.ptr] = []

    def pack_context(self):
        self.ptr -= 1
        return self.stk[self.ptr+1]

    def append_instr(self, ins):
        self.stk[self.ptr] += \
            (ins.splitlines() if isinstance(ins, str) else ins)

    def new_subroutine(self, name, ins):
        self.sub[name] = [f"#{name}:"] + ins

    def assemble(self):
        subs = ["% --- Subroutines ---"]
        for k, v in self.sub.items():
            subs += v
        self.asm = self.stk[0] + subs
        return self.asm


# --------- Instruction Packaging ---------
"""
Here the instructions are packed into functions with this form:

    OPC(r0, r1, ird, ir0, ir1, h, f)

Parameters:
    r0 : operand R0
    r1 : operand R1
    ird: RD inversion flag (0 or 1)
    ir0: R0 inversion flag
    ir1: R1 inversion flag
    h  : HLD flag
    f  : FCT flag

This form is usually used for instructions without operator
correspondence, like CLU, SGN, SMK, etc.

NOTE: LDI is a syntax sugar for LDH and then LDL.
"""


def _apply_invert(rx, irx, typ):
    if isinstance(rx, Register) or isinstance(rx, str):
        return rx, irx
    rx = int(rx)
    if irx == 1:
        rx = -rx if typ == "neg" else ~rx
    return rx, 0

def _to_unsigned(rx):
    return rx & 0xFFFF_FFFF

def _to_signed(rx):
    return -(rx & 0x8000_0000) | (rx & 0xFFFF_FFFF)

def LDI(r0, h=0, f=0):
    return ["LDI", r0, "", 0, 0, 0, h, f]

def LDL(r0, h=0, f=0):
    return ["LDL", r0, "", 0, 0, 0, h, f]

def LDH(r0):
    return ["LDH", r0, "", 0, 0, 0, 0, 0]

def ADD(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "neg")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(r0 + r1)
    else:
        return ["ADD", r0, r1, ird, ir0, ir1, h, f]

def AND(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "inv")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(r0 & r1)
    else:
        return ["AND", r0, r1, ird, ir0, ir1, h, f]

def XOR(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "inv")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(r0 ^ r1)
    else:
        return ["XOR", r0, r1, ird, ir0, ir1, h, f]

def CLU(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "neg")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(-(_to_unsigned(r0) < _to_unsigned(r1)))
    else:
        return ["CLU", r0, r1, ird, ir0, ir1, h, f]

def CLS(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "neg")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(-(_to_signed(r0) < _to_signed(r1)))
    else:
        return ["CLS", r0, r1, ird, ir0, ir1, h, f]

def CEQ(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "neg")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    if isinstance(r0, int) and isinstance(r1, int):
        return LDI(-(_to_unsigned(r0) == _to_unsigned(r1)))
    else:
        return ["CEQ", r0, r1, ird, ir0, ir1, h, f]

def SGN(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "neg")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    if isinstance(r0, int) and isinstance(r1, int):
        r0 = _to_signed(r0)
        r1 = 1 if r1 == 0 else _to_signed(r1)
        return LDI(r0 * r1 // abs(r1))
    else:
        return ["SGN", r0, r1, ird, ir0, ir1, h, f]

def SNE(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    return ["SNE", r0, r1, ird, ir0, ir1, h, f]

def SMK(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "inv")
    return ["SMK", r0, r1, ird, ir0, ir1, h, f]

def MOV(r0, ird=0, ir0=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    return ["MOV", r0, "", ird, ir0, 0, h, f]

def SLL(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    return ["SLL", r0, r1, ird, ir0, ir1, h, f]

def SLA(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    return ["SLA", r0, r1, ird, ir0, ir1, h, f]

def SLC(r0, r1, ird=0, ir0=0, ir1=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    r1, ir1 = _apply_invert(r1, ir1, "neg")
    return ["SLC", r0, r1, ird, ir0, ir1, h, f]

def REV(r0, ird=0, ir0=0, h=0, f=0):
    r0, ir0 = _apply_invert(r0, ir0, "inv")
    return ["REV", r0, "", ird, ir0, 0, h, f]


# --------- Register Management ---------
"""
Registers of a RTMQ core are referenced with:

    R.name

For example: R.TIM, R.ERR, R.R0, etc.
NOTE: register name must be all uppercase.
Assignment to a register:

    R.name = expr

Example:
    R.LED = 0xFF
    R.ERR = ~R.NUL
    R.R0 = R.R1 + 5

NOTE: assignment must be explicit, implicit assignment like these won't work:
    reg = R.R0
    reg = 5

In this case, 'Set' should be used:
    reg = R.R0
    Set(reg, 5)

NOTE:
    1. Signed compare (CLS) is used for '>', '<', '>=' and '<='.
    2. Arithmetic shift (SLA) is used for '>>' and '<<'.
"""


class Register:
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return self.name

    def __lt__(self, other):
        return CLS(self, other)

    def __le__(self, other):
        return CLS(other, self, 1, 0, 0)

    def __gt__(self, other):
        return CLS(other, self)

    def __ge__(self, other):
        return CLS(self, other, 1, 0, 0)

    def __eq__(self, other):
        return CEQ(self, other)

    def __ne__(self, other):
        return CEQ(self, other, 1, 0, 0)

    def __add__(self, other):
        return ADD(self, other)

    def __sub__(self, other):
        return ADD(self, other, 0, 0, 1)

    def __and__(self, other):
        return AND(self, other)

    def __or__(self, other):
        return AND(self, other, 1, 1, 1)

    def __xor__(self, other):
        return XOR(self, other)

    def __invert__(self):
        return MOV(self, 0, 1)

    def __abs__(self):
        return SGN(self, self)

    def __neg__(self):
        return MOV(self, 1, 1)

    def __lshift__(self, other):
        return SLA(self, other)

    def __rshift__(self, other):
        return SLA(self, other, 0, 0, 1)

    def __radd__(self, other):
        return ADD(other, self)

    def __rsub__(self, other):
        return ADD(other, self, 0, 0, 1)

    def __rand__(self, other):
        return AND(other, self)

    def __ror__(self, other):
        return AND(other, self, 1, 1, 1)

    def __rxor__(self, other):
        return XOR(other, self)

    def __rlshift__(self, other):
        return SLA(other, self)

    def __rrshift__(self, other):
        return SLA(other, self, 0, 0, 1)


_AUTO_BUBBLE = contextvars.ContextVar("RTMQ_AUTO_BUBBLE")

class AutoBubble:
    """

    This context is used to indicate whether NOPs be automatically
    inserted after each assignment, i.e., 'bubble' the ALU pipeline.

    The behavior of register assignment expression, 'R.name = expr',
    is defined by this context.

    Parameter: state, can be:
        0: no bubble
        1: bubble

    The default setting for 'ExpFlow' is 1, for 'CfgFlow' is 0.

    NOTE: if deterministic bubbling behavior is required, use function
    'Set' instead.

    """
    def __init__(self, state):
        self.state = state

    def __enter__(self):
        self.token = _AUTO_BUBBLE.set(self.state)

    def __exit__(self, exc_type, exc_val, exc_tb):
        _AUTO_BUBBLE.reset(self.token)


def _prep_operand(rx, trx, core):
    if (isinstance(rx, int) and abs(rx) > 63) or \
       (isinstance(rx, str) and rx != ""):
        ins = [f"LDH - -  {trx}  {rx}", f"LDL - -  {trx}  {rx}"]
        bub = StdCore.N_PLI
        rx = trx
    elif isinstance(rx, Register) and core.regs[str(rx)] > 63:
        ins = [f"MOV - -  {trx}  {rx}"]
        bub = StdCore.N_PLA
        rx = trx
    else:
        ins = []
        bub = 0
        rx = str(rx)
    return ins, rx, bub


def _assign_reg(rd, expr, core, bubble=None):
    if isinstance(expr, Register):
        expr = MOV(expr)
    elif not isinstance(expr, list):
        expr = LDI(int(expr))
    bubble = _AUTO_BUBBLE.get() if bubble is None else bubble
    try:
        opc, r0, r1, ird, ir0, ir1, h, f = expr
        h = ("-", "H")[h]
        f = ("-", "F")[f]
        ins = []
        if opc in {"LDI", "LDH", "LDL"}:
            if opc != "LDL":
                ins += [f"LDH - -  {rd}  {r0}"]
            if opc != "LDH":
                ins += [f"LDL {h} {f}  {rd}  {r0}"]
            if bubble:
                ins += ["NOP - -"] * StdCore.N_PLI
            return ins

        in0, r0, bu0 = _prep_operand(r0, "TR0", core)
        in1, r1, bu1 = _prep_operand(r1, "TR1", core)
        if bu1 > bu0:
            ins = in1 + in0
            bub = bu1 - len(in0)
        elif bu0 > bu1:
            ins = in0 + in1
            bub = bu0 - len(in1)
        else:
            ins = in0 + in1
            bub = bu1
        ins += ["NOP - -"] * bub
        ird = (" ", "!")[ird]
        ir0 = (" ", "!")[ir0]
        ir1 = (" ", "!")[ir1]
        ins += [f"{opc} {h} {f} {ird}{rd} {ir0}{r0} {ir1}{r1}"]
        if bubble:
            ins += ["NOP - -"] * StdCore.N_PLA
        return ins
    except:
        raise RuntimeError(f"Invalid expression: {expr}.")


def Set(rd, expr, bubble=0):
    """
    
    Implement RTMQ register assignment operation.
    'Set(R.rd, expr)' is equivalent to 'R.rd = expr',
    other than that 'Set' is able to determine
    bubbling behavior explicitly.

    Parameters
    ----------
    rd : <Register>
        Destination register RD.
    expr : integer, <Register> OR 1-step operation involving Registers
        Expression to be assigned to RD.
    bubble : integer, optional
        0: no bubble; 
        1: bubble;
        The default is 0.

    Returns
    -------
    None.

    """
    asmb = _get_asmb()
    asmb.append_instr(_assign_reg(rd, expr, asmb.core, bubble))


class RegMgr():
    def __setattr__(self, name, value):
        if name.isupper():
            Set(name, value, None)
        else:
            self.__dict__[name] = value

    def __getattr__(self, name):
        if name in self.__dict__.keys():
            return self.__dict__[name]
        else:
            return Register(name)

R = RegMgr()


# -------- Flow Control Structures --------


def ASM(ins):
    _get_asmb().append_instr(ins)


def Pass(n=1):
    ASM(["NOP - -"] * n)


def Hold():
    ASM("NOP H -")


def HardWait(n):
    cyc = StdCore.N_PLA + StdCore.N_PLM + 2
    if n <= cyc + 3:
        Pass(n)
    else:
        n = n - 3
        Set(R.TR0, 1 - n // cyc, 1)
        Set(R.PTR, SNE(R.PTR, R.TR0))
        Set(R.TR0, R.TR0 + 1)
        Pass(cyc + n % cyc - 2)


def Label(label):
    ASM(f"#{label}:")


def Jump(dst):
    if isinstance(dst, Register):
        ASM(f"MOV - F  PTR  {dst}")
    else:
        if isinstance(dst, str):
            dst = f"#{dst}"
        ASM(f"LDL - F  PTR  {dst}")


def JumpIf(cond, label, inv=0):
    if isinstance(cond, Register):
        Set(R.TR0, LDI(f"#{label}"), 1)
        Set(R.PTR, SNE(R.TR0, cond, ird=inv, f=1))
    else:
        cond[3] = cond[3] ^ inv
        Set(R.TR1, cond)
        Set(R.TR0, LDI(f"#{label}"))
        Pass(StdCore.N_PLA - 2)
        Set(R.PTR, SNE(R.TR0, R.TR1, f=1))


def StartTimedTask(time):
    if isinstance(time, int):
        ASM([f"LDH - -  TR0  {time}",
             f"LDL H -  TR0  {time}"])
        time = R.TR0
    elif isinstance(time, list):
        Set(R.TR0, time)
        Pass(StdCore.N_PLA - 2)
        Hold()
        time = R.TR0
    else:
        Hold()
    Pass()
    Set(R.TIM, time)
    Set(R.NOW, R.WCK + time)


def Call(sub):
    ASM([ "ADD - -  RTS  PTR  2",
         f"LDL - F  PTR  #{sub}"])


def Return():
    ASM("MOV - F  PTR  RTS")


class LoopForever:
    def __init__(self):
        self.asmb = _get_asmb()
        if not hasattr(self.asmb, "cnt_lfrvr"):
            self.asmb.cnt_lfrvr = 0
        self.tag = f"LoopForever_{self.asmb.cnt_lfrvr}"
        self.asmb.cnt_lfrvr += 1

    def __enter__(self):
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        Jump(self.tag)
        s = self.asmb.pack_context()
        Label(self.tag)
        self.asmb.append_instr(s)


class LoopWhile:
    def __init__(self, cond):
        self.asmb = _get_asmb()
        if not hasattr(self.asmb, "cnt_lwh"):
            self.asmb.cnt_lwh = 0
        self.tag = f"LoopWhile_{self.asmb.cnt_lwh}"
        self.asmb.cnt_lwh += 1
        self.tb = self.tag + "_Begin"
        self.te = self.tag + "_End"
        self.cond = cond

    def __enter__(self):
        Label(self.tb)
        JumpIf(self.cond, self.te, inv=1)
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        Jump(self.tb)
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.te)


class Scan:
    def __init__(self, reg, ini, fin, stp):
        self.asmb = _get_asmb()
        if not hasattr(self.asmb, "cnt_scan"):
            self.asmb.cnt_scan = 0
        self.tag = f"Scan_{self.asmb.cnt_scan}"
        self.asmb.cnt_scan += 1
        self.tb = self.tag + "_Begin"
        self.te = self.tag + "_End"
        self.r = reg
        self.i = ini if isinstance(ini, Register) else int(ini)
        self.f = fin if isinstance(fin, Register) else int(fin)
        self.s = stp if isinstance(stp, Register) else int(stp)

    def __enter__(self):
        Set(self.r, self.i, 1)
        Label(self.tb)
        if isinstance(self.s, int) and isinstance(self.f, int):
            Set(R.TR1, SGN(self.r, self.s))
            Set(R.TR0, SGN(self.f, self.s))
            Pass(StdCore.N_PLA - 2)
        else:
            st = self.s
            if (isinstance(st, int) and abs(st) >= 64) or \
               (isinstance(st, Register) and \
                self.asmb.core.regs[str(st)] > 63):
                st = R.TR1
            Set(R.TR0, SGN(self.f, self.s))
            Set(R.TR1, SGN(self.r, st), 1)
        JumpIf(R.TR1 >= R.TR0, self.te)
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        Set(self.r, self.r + self.s)
        Jump(self.tb)
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.te)


class Repeat():
    def __init__(self, reg, rep):
        self.asmb = _get_asmb()
        if not hasattr(self.asmb, "cnt_rep"):
            self.asmb.cnt_rep = 0
        self.tag = f"Repeat_{self.asmb.cnt_rep}"
        self.asmb.cnt_rep += 1
        self.tb = self.tag + "_Begin"
        self.te = self.tag + "_End"
        self.r = reg
        self.p = rep

    def __enter__(self):
        Set(self.r, 0, 1)
        Label(self.tb)
        self.asmb.new_context()
        JumpIf(self.r >= self.p, self.te)

    def __exit__(self, exc_type, exc_val, exc_tb):
        Set(self.r, self.r + 1)
        Jump(self.tb)
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.te)


class Subroutine:
    def __init__(self, tag, protect=[]):
        self.tag = tag
        self.prt = protect
        self.asmb = _get_asmb()

    def __enter__(self):
        self.asmb.new_context()
        if self.prt != []:
            for r in self.prt:
                Set(R.DTS, r)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.prt != []:
            self.prt.reverse()
            for r in self.prt:
                Set(r, R.DTS, 1)
        Return()
        self.asmb.new_subroutine(self.tag, self.asmb.pack_context())


class If:
    def __init__(self, cond):
        self.asmb = _get_asmb()
        a = self.asmb
        if not hasattr(a, "cnt_if"):
            a.cnt_if = 0
        self.tag_els = f"If_Else_{a.cnt_if}"
        self.tag_ret = f"If_Retn_{a.cnt_if}"
        a.cnt_if += 1
        self.cond = cond

    def __enter__(self):
        JumpIf(self.cond, self.tag_els, inv=1)
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        Jump(self.tag_ret)
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.tag_els)
        Label(self.tag_ret)


class Elif:
    def __init__(self, cond):
        self.asmb = _get_asmb()
        a = self.asmb
        if not hasattr(a, "cnt_elif"):
            a.cnt_elif = 0
        self.tag_els = f"Elif_Else_{a.cnt_elif}"
        a.cnt_elif += 1
        self.cond = cond
        if not a.stk[a.ptr][-1].startswith("#If_Retn_"):
            raise RuntimeError("Elif should be immediately after If or Elif.")
        self.tag_ret = a.stk[a.ptr][-1][1:-1]
        del a.stk[a.ptr][-1]

    def __enter__(self):
        JumpIf(self.cond, self.tag_els, inv=1)
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        Jump(self.tag_ret)
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.tag_els)
        Label(self.tag_ret)


class Else:
    def __init__(self):
        self.asmb = _get_asmb()
        a = self.asmb
        if not a.stk[a.ptr][-1].startswith("#If_Retn_"):
            raise RuntimeError("Else should be immediately after If or Elif.")
        self.tag_ret = a.stk[a.ptr][-1][1:-1]
        del a.stk[a.ptr][-1]

    def __enter__(self):
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        Label(self.tag_ret)


class UntimedTask:
    def __init__(self):
        self.asmb = _get_asmb()

    def __enter__(self):
        Hold()
        self.asmb.new_context()

    def __exit__(self, exc_type, exc_val, exc_tb):
        s = self.asmb.pack_context()
        self.asmb.append_instr(s)
        ASM("MOV - -  TRG  TRG")


_ASMB = contextvars.ContextVar("RTMQ_ASMB")

def _get_asmb():
    try:
        return _ASMB.get()
    except LookupError:
        raise RuntimeError("Flow generation funcions cannot be called " + \
                           "outside of <ExpFlow> or " + \
                           "<CfgFlow> context.")


class ExpFlow:
    def __init__(self, core):
        self.token = _ASMB.set(Compiler(core))
        self.typ = "exp"
        self.core = core

    def __enter__(self):
        _AUTO_BUBBLE.set(1)
        ASM("MOV - -  TRG  TRG")
        ASM("MOV - -  ERR  NUL")

    def __exit__(self, exc_type, exc_val, exc_tb):
        ASM(["NOP H -", "NOP H -"])
        self.assem = _get_asmb().assemble()
        _ASMB.reset(self.token)


class CfgFlow:
    def __init__(self, core):
        self.token = _ASMB.set(Compiler(core))
        self.typ = "cfg"
        self.core = core

    def __enter__(self):
        _AUTO_BUBBLE.set(0)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.assem = _get_asmb().assemble()
        _ASMB.reset(self.token)


# -------- System Peripheral Functions --------


def Wait(time):
    if isinstance(time, int):
        Set(R.TIM, LDI(time, h=1))
    else:
        Pass()
        Set(R.TIM, MOV(time, h=1))


def WaitTrigger(trig):
    Hold()
    Set(R.TRG, trig)


def SerialSend(dat, wait=True):
    if wait:
        WaitTrigger(1)
    Set(R.RTD0, dat)


def DatCacheRead(addr):
    Set(R.DCA, addr, 1)
    Pass(StdCore.N_PLM)


def DatCacheWrite(addr, dat):
    Set(R.DCA, addr)
    Set(R.DCD, dat)


def InsCacheRead(addr):
    Set(R.ICA, addr, 1)
    Pass(StdCore.N_PLM)


def InsCacheWrite(addr, dat):
    Set(R.ICA, addr)
    Set(R.ICD, dat)


def ResetWallClock(ofs):
    Set(R.WCK, ofs)


def ReadWallClock(h, l):
    Set(l, R.WCK)
    Set(h, R.TIM, 1)


# -------- UART Interface (TO BE REMOVED) --------


import cion.core.hardware.rtmq.RTMQ_Assembler as asmblr
import serial

class UARTDevice:
    def __init__(self, port, baud):
        self.port = port if isinstance(port, list) else [port]
        self.baud = baud
        self.ser = None
        self.port_cnt = 0

    def open_port(self, tout=5):
        if self.ser is None:
            self.ser = []
            for p in self.port:
                s = serial.Serial(p, self.baud, timeout=tout)
                s.stopbits = serial.STOPBITS_ONE
                self.ser += [s]
        self.port_cnt += 1

    def close_port(self):
        self.port_cnt -= 1
        if self.port_cnt == 0:
            for s in self.ser:
                s.close()
            self.ser = None

    def read(self, port, ret_cnt=None):
        s = self.ser[port]
        if ret_cnt is None:
            ret_cnt = s.in_waiting // 5
        return asmblr.get_payloads(s.read(ret_cnt * 5))

    def write_stream(self, port, stream):
        self.ser[port].write(stream)
    
    def simulate(self, fn, flow):
        prg = asmblr.asm_program(flow.assem, flow.core)
        with open(fn, "w") as f:
            f.write(asmblr.sim_form(prg))

    def run(self, port, addr, flow, ret_cnt=0):
        buf = asmblr.reach_node(addr)
        if flow.typ == "cfg":
            buf += asmblr.suspend_node(addr) + \
                   asmblr.node_exec(addr, flow.assem, flow.core) + \
                   asmblr.resume_node(addr)
        else:
            prg = asmblr.asm_program(flow.assem, flow.core)
            buf += asmblr.suspend_node(addr) + \
                   asmblr.download_to(addr, prg) + \
                   asmblr.reset_node(addr) + \
                   asmblr.resume_node(addr)
        self.open_port(1)
        self.write_stream(port, buf)
        if ret_cnt > 0:
            ret = self.read(port, ret_cnt)
        self.close_port()
        if ret_cnt > 0:
            return ret
